from mpl_toolkits import mplot3d
import numpy as np
import matplotlib.pyplot as plt

# fig = plt.figure()
fig, (ax1, ax2) = plt.subplots(2, 1, gridspec_kw={'height_ratios': [1,1]})


file="pydata.txt"

mnid = np.array([int(x.split(' ')[0]) for x in open(file).readlines()])
kp_x = np.array([float(x.split(' ')[1]) for x in open(file).readlines()])
kp_y = np.array([float(x.split(' ')[2]) for x in open(file).readlines()])
mp_x = np.array([float(x.split(' ')[3]) for x in open(file).readlines()])
mp_y = np.array([float(x.split(' ')[4]) for x in open(file).readlines()])
mp_z = np.array([float(x.split(' ')[5]) for x in open(file).readlines()])
ct_x = np.array([float(x.split(' ')[6]) for x in open(file).readlines()])
ct_y = np.array([float(x.split(' ')[7]) for x in open(file).readlines()])
ct_z = np.array([float(x.split(' ')[8]) for x in open(file).readlines()])

alldata = np.stack([mnid, kp_x, kp_y,mp_x,mp_y,mp_z,ct_x,ct_y,ct_z],axis=1)
# print(alldata.shape)
# print(alldata[0,:])

frameWiseData = np.split(alldata, np.where(np.diff(alldata[:,0]))[0]+1)
# print(len(frameWiseData)) # number of frames
# print(frameWiseData[0].shape) # number of points in frame 0
# print(frameWiseData[0]) # points in frame 0


# print(frameWiseData[0][:,3:6])
# print(frameWiseData[0][:,6:9])

# dist = np.linalg.norm(frameWiseData[0][:,3:6] - frameWiseData[0][:,6:9], axis=1) 
# print(dist.shape)
# print(dist)

# ax.scatter(frameWiseData[0][:,1],frameWiseData[0][:,2],s=20,c=dist, marker = 'o', cmap = "hot" );
# plt.show()


for frame in frameWiseData:
	id = str(int(frame[0][0]))
	dist = np.linalg.norm(frame[:,3:6] - frame[:,6:9], axis=1) 
	dist = (dist*100) / np.linalg.norm(dist) # normalize to give uniform color shade

	# ax1 = fig.add_subplot(211)	
	ax1.scatter(frame[:,1],frame[:,2],s=20,c=dist, marker = 'o', cmap = "hot" ); # https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
	ax1.axis('off')

	# ax2 = fig.add_subplot(212)	
	im = plt.imread("../Downloads/orb_custom_data/custom_data6/out/rgb/frame"+id+".jpg")
	ax2.imshow(im)
	ax2.axis('off')

	plt.savefig("data/frame"+id+'_.png',bbox_inches='tight') 
	ax1.clear()
	ax2.clear()


# for frame in frameWiseData:
# 	id = str(int(frame[0][0]))
# 	im = plt.imread("../Downloads/orb_custom_data/custom_data6/out/rgb/frame"+id+".jpg")
# 	plt.scatter([10], [20])
# 	# plt.axis('off')
# 	plt.savefig(id+'.png',bbox_inches='tight')
